
#include "gcd_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include <algorithm>

namespace optiling {
    constexpr uint32_t BLOCK_SIZE = 512;
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{
    GcdTilingData tiling;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());

    const gert::StorageShape* x1_StorageShape = context->GetInputShape(0);
    const gert::StorageShape* x2_StorageShape = context->GetInputShape(1);

    uint32_t x1_size = x1_StorageShape->GetStorageShape().GetShapeSize();
    uint32_t x2_size = x2_StorageShape->GetStorageShape().GetShapeSize();
    uint32_t y_size;
    uint32_t BLOCK_DIM = 1;
    uint32_t inputBytes = GetSizeByDataType(context->GetInputDesc(0)->GetDataType());
    if(x1_size != x2_size)
    {
        uint32_t x1_dim = x1_StorageShape->GetStorageShape().GetDimNum();
        uint32_t x2_dim = x2_StorageShape->GetStorageShape().GetDimNum();
        uint32_t y_dim;
        uint32_t x1_shape[5];
        uint32_t x2_shape[5];
        uint32_t y_shape[5];
        y_size = 1;
        y_dim = x1_dim;
        if(y_dim < x2_dim) y_dim = x2_dim;

        uint32_t x1_start = y_dim - x1_dim;
        uint32_t x2_start = y_dim - x2_dim;

        for(int i=0; i< y_dim; i++)
        {
            if(i < x1_start) x1_shape[i] = 1;
            else x1_shape[i] = x1_StorageShape->GetStorageShape().GetDim(i - x1_start);

            if(i < x2_start) x2_shape[i] = 1;
            else x2_shape[i] = x2_StorageShape->GetStorageShape().GetDim(i - x2_start);

            y_shape[i] = x1_shape[i];
            if(y_shape[i] < x2_shape[i]) y_shape[i] = x2_shape[i];

            y_size *= y_shape[i];
        }
        
        tiling.set_y_dim(y_dim);
        tiling.set_x1_shape(x1_shape);
        tiling.set_x2_shape(x2_shape);
        tiling.set_y_shape(y_shape);

        context->SetTilingKey(3);
    }
    else if(inputBytes == 4)
    {
         context->SetTilingKey(1);
        
        if(y_size >= 40*16)
        {
            BLOCK_DIM = ascendcPlatform.GetCoreNum();
        }
        uint32_t inputLength = inputBytes * y_size;
        uint32_t totalLengthAligned = ((inputLength + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE;
        uint32_t formerNum = (totalLengthAligned / BLOCK_SIZE) % BLOCK_DIM;
        uint32_t tailLength = (totalLengthAligned / BLOCK_DIM / BLOCK_SIZE) * BLOCK_SIZE;
        tiling.set_formerNum(formerNum);
        tiling.set_tailLength(tailLength/inputBytes);
    }
    else
    {
        y_size = x1_size;
        context->SetTilingKey(2);
    }

    tiling.set_size(y_size);

    context->SetBlockDim(BLOCK_DIM);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

  return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class Gcd : public OpDef {
public:
    explicit Gcd(const char* name) : OpDef(name)
    {
        this->Input("x1")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT16, ge::DT_INT32, ge::DT_INT64})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("x2")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT16, ge::DT_INT32, ge::DT_INT64})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT16, ge::DT_INT32, ge::DT_INT64})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");

    }
};

OP_ADD(Gcd);
}
